/*
** Copyright (C) 2010 Calin Radoni
**
** This file is part of CHScanner.
**
** CHScanner is free software: you can redistribute it and/or modify
** it under the terms of the GNU General Public License as published by
** the Free Software Foundation, either version 3 of the License, or
** (at your option) any later version.
**
** CHScanner is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
** GNU General Public License for more details.
**
** You should have received a copy of the GNU General Public License
** along with CHScanner. If not, see <http://www.gnu.org/licenses/>.
*/

#include "StdAfx.h"
#include ".\scanmethodsnmp.h"

//--------------------------------------------------------------------

ScanMethodSnmp *snmpMethod;

//--------------------------------------------------------------------

ScanMethodSnmp::ScanMethodSnmp(void)
{
	snmpMethod = NULL;

	myType = __wmtWinSnmp;
	communityName.Empty();
}

ScanMethodSnmp::~ScanMethodSnmp(void)
{
}

//--------------------------------------------------------------------

int ScanMethodSnmp::InternalInit(void)
{
	if(workPhase!=__wptScan) return MSR_ERR_PHASE;

	communityName = scanParameters.Trim();

	return MSR_ERR_OK;
}

//--------------------------------------------------------------------

int ScanMethodSnmp::InternalExecute(void)
{
	ScanPacket* pack;

	if(addrList==NULL) return MSR_ERR_PARAMETER;

	if(!winSnmp.SnmpInit()) return MSR_ERR_FAILED;

	snmpMethod = this;

	in_addr ia;
	CString strLocalIP;
	HSNMP_SESSION hSession;

	bool done;
	MSnmpSysRec snmpSystemInfo;

	ia.S_un.S_addr = winPcap->GetLocalIP();
	strLocalIP = inet_ntoa(ia);

	ScanPair sp;
	while(GetNextPair(&sp, __sptIpOnly) && !stopStopStop){
		if(sp.addr!=NULL){
			if(!sp.addr->addrIP.IsEmpty()){
				if(sp.addr->alive){
					if(config->xMinDelay>0){
						if(lastScanTime!=0){
							waitDone = false;
							while(!waitDone){
								CheckResponses();
								SleepEx(1, TRUE);
								scanTime = GetTickCount();
								scanTime = ComputeDeltaTime(lastScanTime, scanTime);
								if(scanTime>=config->xMinDelay)
									waitDone = true;
							}
						}
						lastScanTime = GetTickCount();
					}

					pack = FindFirstUnusedPacket();
					while(pack==NULL){
						SleepEx(waitTime, TRUE);
						CheckResponses();
						CheckDeaths();
						pack = FindFirstUnusedPacket();
					}

					if(pack!=NULL){
						pack->Initialize(true);
						pack->sTime       = GetTickCount();
						pack->addr        = sp.addr->addrIP.GetD();
						pack->port        = 161;
						pack->methodType  = myType;
						pack->fullAddress = sp.addr;
						pack->machineInfo = FindMachineByIp(pack->addr);

						if(winSnmp.SnmpInitSession(SnmpCallbackFunction, (void*)pack)){
							hSession = winSnmp.GetCurrentSessionHandle();
							pack->paramVP = hSession;
							if(!winSnmp.SendSystemQuery(strLocalIP, sp.addr->addrIP.GetS(), communityName)){
								pack->Initialize(false);
								winSnmp.SnmpDestroySession(hSession);
							}
							else{
								done = false;
								while(!done){
									if(!winSnmp.RecvSystemQuery(hSession)){
										SleepEx(1, TRUE);
										CheckDeaths();
										if(pack->paramVP==NULL) done = true;
									}
									else{
										snmpSystemInfo.sysDescription = winSnmp.xSystemQueryResult.sysDescription;
										snmpSystemInfo.sysVendorOid   = winSnmp.xSystemQueryResult.sysVendorOid;
										snmpSystemInfo.sysUpTime      = winSnmp.xSystemQueryResult.sysUpTime;
										snmpSystemInfo.sysContact     = winSnmp.xSystemQueryResult.sysContact;
										snmpSystemInfo.sysName        = winSnmp.xSystemQueryResult.sysName;
										snmpSystemInfo.sysLocation    = winSnmp.xSystemQueryResult.sysLocation;
										snmpSystemInfo.sysServices    = winSnmp.xSystemQueryResult.sysServices;
										if(pack->machineInfo!=NULL){
											pack->machineInfo->snmpSysRec.LoadFromPointer(&snmpSystemInfo);
											pack->machineInfo->address.alive = true;
										}
										done = true;
									}
								}
								winSnmp.SnmpDestroySession(hSession);
								packSent++;
								if((packSent % notifNetBios) == 0 && packSent>0)
									if(hWnd!=NULL)
										PostMessage(hWnd, CHS_MSG_PACKETS_SENT, 0, (LPARAM)packSent);
							}
						}
						else{
							pack->Initialize(false);
						}
					}
				}
			}
		}
	}

	// wait until the end
	while(PacketsUsed()>0){
		CheckResponses();
		CheckDeaths();
		SleepEx(1, TRUE);
	}

	snmpMethod = NULL;
	winSnmp.SnmpDestroy();

	return MSR_ERR_OK;
}

//--------------------------------------------------------------------

void ScanMethodSnmp::SnmpReceiveFunction(HSNMP_SESSION hSessionIn, WPARAM wParam, LPARAM lParam, LPVOID lpClientData)
{
/*
	MMachineInfo *machineInfo;
	MSnmpSysRec snmpSystemInfo;
	ScanPacket* pack;
	HSNMP_SESSION hSession;
*/
	return;
/*
	if(wParam==0){
		if(winSnmp.RecvSystemQuery(hSessionIn)){
			pack = (ScanPacket*)lpClientData;
			if(pack!=NULL){
				hSession = (HSNMP_SESSION)pack->paramVP;
				if(hSession==hSessionIn){
					try{
						snmpSystemInfo.sysDescription = winSnmp.xSystemQueryResult.sysDescription;
						snmpSystemInfo.sysVendorOid   = winSnmp.xSystemQueryResult.sysVendorOid;
						snmpSystemInfo.sysUpTime      = winSnmp.xSystemQueryResult.sysUpTime;
						snmpSystemInfo.sysContact     = winSnmp.xSystemQueryResult.sysContact;
						snmpSystemInfo.sysName        = winSnmp.xSystemQueryResult.sysName;
						snmpSystemInfo.sysLocation    = winSnmp.xSystemQueryResult.sysLocation;
						snmpSystemInfo.sysServices    = winSnmp.xSystemQueryResult.sysServices;
					}
					catch(...){
					}

					machineInfo = pack->machineInfo;
					if(machineInfo!=NULL){
						machineInfo->snmpSysRec.LoadFromPointer(&snmpSystemInfo);
						machineInfo->address.alive = true;
					}
				}
			}
		}
	}

	if(pack!=NULL)
		pack->Initialize(false);
*/
}

//--------------------------------------------------------------------

SNMPAPI_STATUS CALLBACK SnmpCallbackFunction(HSNMP_SESSION hSessionIn, HWND hWnd, UINT wMsg, WPARAM wParam, LPARAM lParam, LPVOID lpClientData)
{
	if(snmpMethod!=NULL)
		snmpMethod->SnmpReceiveFunction(hSessionIn, wParam, lParam, lpClientData);

	return SNMPAPI_SUCCESS;
}

//--------------------------------------------------------------------

int ScanMethodSnmp::InternalDeath(ScanPacket* spIn)
{
	HSNMP_SESSION hSession;

	hSession = (HSNMP_SESSION)spIn->paramVP;
	winSnmp.SnmpDestroySession(hSession);

/*
	CString str;
	str.Format(_T("---::D %s:%d\n"), spIn->fullAddress->addrIP.GetS(), spIn->port);
	OutputDebugString(str);
*/

	return MSR_ERR_OK;
}

//---------------------------------------------------------------------------

void ScanMethodSnmp::ProcessPortUnreacheable(BYTE* dataIn, int sizeIn)
{
	in_addr      inAddr;
	ScanPacket   *pack;
	IPHeaderPG   *iph;
	UDPHeaderPG  *udph;
	int          dataOffset;

	iph = (IPHeaderPG*)dataIn;
	dataOffset  = (int)(iph->HLV & 0x0F);
	dataOffset *= 4;
	if(sizeIn<(dataOffset + 8)){
		// incomplete packet
		return;
	}

	udph = (UDPHeaderPG*)&dataIn[dataOffset];

	inAddr.S_un.S_un_b.s_b1 = iph->Destination[0];
	inAddr.S_un.S_un_b.s_b2 = iph->Destination[1];
	inAddr.S_un.S_un_b.s_b3 = iph->Destination[2];
	inAddr.S_un.S_un_b.s_b4 = iph->Destination[3];

	pack = FindPacket(inAddr.S_un.S_addr, ntohs(udph->DestinationPort), myType);
	if(pack==NULL) return;

	HSNMP_SESSION hSession;

	hSession = (HSNMP_SESSION)pack->paramVP;
	winSnmp.SnmpDestroySession(hSession);

/*
	CString str;
	str.Format(_T("---::P %s:%d\n"), pack->fullAddress->addrIP.GetS(), pack->port);
	OutputDebugString(str);
*/
	pack->Initialize(false);
}

//---------------------------------------------------------------------------
